package com.dx.fedpauth.common.filter;

import com.vdurmont.emoji.EmojiParser;
import org.apache.commons.lang3.StringUtils;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.safety.Whitelist;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

public class XssRequestWrapper extends HttpServletRequestWrapper {

    /**
     * 代码块的分隔符
     */
    private static final String TRIPLE_DELIMITER = "```";
    private static final String SINGLE_DELIMITER = "`";

    XssRequestWrapper(HttpServletRequest servletRequest) {
        super(servletRequest);
    }

    @Override
    public String[] getParameterValues(String parameter) {
        String[] values = super.getParameterValues(parameter);

        if (values == null) {
            return null;
        }

        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = stripXSS(values[i]);
        }

        return encodedValues;
    }

    @Override
    public String getParameter(String parameter) {
        String value = super.getParameter(parameter);

        return stripXSS(value);
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        return stripXSS(value);
    }

    public static String stripXSS(String value) {
        if (value == null) {
            return null;
        }

        value = EmojiParser.removeAllEmojis(value);

        String result = onePass(value, TRIPLE_DELIMITER);

        // 将小于号和大于号转换回去
        result = result.replaceAll("&lt;", "<").replaceAll("&gt;", ">");

        return result;
    }

    private static String onePass(String value, String delimiter) {
        // 过滤器的输出设置
        Document.OutputSettings outputSettings = new Document.OutputSettings().prettyPrint(false);

        // 采用一个指针遍历整个字符串，代码块内部的内容不进行过滤，不在代码块内部的内容要过滤
        // 当前遍历的部分
        String currentBlock;
        // 尚未遍历的部分
        String remainingPart = value;
        // 过滤之后的结果
        String result = "";

        // 指示当前指针是否位于代码块内部
        boolean signalIn = false;

        // 到整个字符串遍历完为止
        while (StringUtils.length(remainingPart) > 0) {
            currentBlock = StringUtils.substringBefore(remainingPart, delimiter);
            // 不在代码块内部的时候进行过滤
            if (StringUtils.isNotEmpty(currentBlock) && !signalIn) {
                // 先对多行代码块（三个`符号包围的）进行一轮过滤
                if (delimiter.equals(TRIPLE_DELIMITER)) {
                    // 在多行代码块之外的地方再针对单行代码块进行一轮过滤，递归调用过滤函数本身
                    currentBlock = onePass(currentBlock, SINGLE_DELIMITER);
                } else {
                    // 此处为多行代码块以及单行代码块之外的地方，调用Jsoup清理XSS攻击
                    currentBlock = Jsoup.clean(currentBlock, "",
                            Whitelist.relaxed().addProtocols("img","src","data"), outputSettings);
                }
            }
            result += currentBlock;

            if (StringUtils.contains(remainingPart, delimiter)) {
                remainingPart = StringUtils.substringAfter(remainingPart, delimiter);
                result += delimiter;
                // 确保分隔符是成双出现的情况下，才认为是在代码块内
                if (StringUtils.contains(remainingPart, delimiter) || signalIn) {
                    signalIn = !signalIn;
                }
            } else {
                remainingPart = "";
            }
        }
        return result;
    }

}